// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/json/converter.h>

#include <memory>
#include <string_view>
#include <utility>

#include <nebula/core/array.h>
#include <nebula/array/builder_binary.h>
#include <nebula/array/builder_decimal.h>
#include <nebula/array/builder_primitive.h>
#include <nebula/array/builder_time.h>
#include <nebula/json/parser.h>
#include <nebula/types/type.h>
#include <turbo/base/checked_cast.h>
#include <nebula/numeric/decimal.h>
#include <turbo/log/logging.h>
#include <nebula/util/value_parsing.h>


namespace nebula::json {

    template<typename... Args>
    turbo::Status GenericConversionError(const DataType &type, Args &&... args) {
        return turbo::invalid_argument_error("Failed to convert JSON to ", type, std::forward<Args>(args)...);
    }

    namespace {

        const DictionaryArray &GetDictionaryArray(const std::shared_ptr<Array> &in) {
                    DKCHECK_EQ(in->type_id(), Type::DICTIONARY);
            auto dict_type = turbo::checked_cast<const DictionaryType *>(in->type().get());
                    DKCHECK_EQ(dict_type->index_type()->id(), Type::INT32);
                    DKCHECK_EQ(dict_type->get_value_type()->id(), Type::STRING);
            return turbo::checked_cast<const DictionaryArray &>(*in);
        }

        template<typename ValidVisitor, typename NullVisitor>
        turbo::Status VisitDictionaryEntries(const DictionaryArray &dict_array,
                                      ValidVisitor &&visit_valid, NullVisitor &&visit_null) {
            const StringArray &dict = turbo::checked_cast<const StringArray &>(*dict_array.dictionary());
            const Int32Array &indices = turbo::checked_cast<const Int32Array &>(*dict_array.indices());
            for (int64_t i = 0; i < indices.length(); ++i) {
                if (indices.is_valid(i)) {
                    TURBO_RETURN_NOT_OK(visit_valid(dict.get_view(indices.get_view(i))));
                } else {
                    TURBO_RETURN_NOT_OK(visit_null());
                }
            }
            return turbo::OkStatus();
        }

    }  // namespace

// base class for types which accept and output non-nested types
    class PrimitiveConverter : public Converter {
    public:
        PrimitiveConverter(MemoryPool *pool, std::shared_ptr<DataType> out_type)
                : Converter(pool, out_type) {}
    };

    class NullConverter : public PrimitiveConverter {
    public:
        using PrimitiveConverter::PrimitiveConverter;

        turbo::Status Convert(const std::shared_ptr<Array> &in, std::shared_ptr<Array> *out) override {
            if (in->type_id() != Type::NA) {
                return GenericConversionError(*out_type_, " from ", *in->type());
            }
            *out = in;
            return turbo::OkStatus();
        }
    };

    class BooleanConverter : public PrimitiveConverter {
    public:
        using PrimitiveConverter::PrimitiveConverter;

        turbo::Status Convert(const std::shared_ptr<Array> &in, std::shared_ptr<Array> *out) override {
            if (in->type_id() == Type::NA) {
                return MakeArrayOfNull(boolean(), in->length(), pool_).try_value(out);
            }
            if (in->type_id() != Type::BOOL) {
                return GenericConversionError(*out_type_, " from ", *in->type());
            }
            *out = in;
            return turbo::OkStatus();
        }
    };

    template<typename T>
    class NumericConverter : public PrimitiveConverter {
    public:
        using value_type = typename T::c_type;

        NumericConverter(MemoryPool *pool, const std::shared_ptr<DataType> &type)
                : PrimitiveConverter(pool, type), numeric_type_(turbo::checked_cast<const T &>(*type)) {}

        turbo::Status Convert(const std::shared_ptr<Array> &in, std::shared_ptr<Array> *out) override {
            if (in->type_id() == Type::NA) {
                return MakeArrayOfNull(out_type_, in->length(), pool_).try_value(out);
            }
            const auto &dict_array = GetDictionaryArray(in);

            using Builder = typename TypeTraits<T>::BuilderType;
            Builder builder(out_type_, pool_);
            TURBO_RETURN_NOT_OK(builder.resize(dict_array.indices()->length()));

            auto visit_valid = [&](std::string_view repr) {
                value_type value;
                if (!nebula::internal::ParseValue(numeric_type_, repr.data(), repr.size(), &value)) {
                    return GenericConversionError(*out_type_, ", couldn't parse:", repr);
                }

                builder.unsafe_append(value);
                return turbo::OkStatus();
            };

            auto visit_null = [&]() {
                builder.unsafe_append_null();
                return turbo::OkStatus();
            };

            TURBO_RETURN_NOT_OK(VisitDictionaryEntries(dict_array, visit_valid, visit_null));
            return builder.finish(out);
        }

        const T &numeric_type_;
    };

    template<typename T>
    class DecimalConverter : public PrimitiveConverter {
    public:
        using value_type = typename TypeTraits<T>::BuilderType::value_type;

        DecimalConverter(MemoryPool *pool, const std::shared_ptr<DataType> &type)
                : PrimitiveConverter(pool, type) {}

        turbo::Status Convert(const std::shared_ptr<Array> &in, std::shared_ptr<Array> *out) override {
            if (in->type_id() == Type::NA) {
                return MakeArrayOfNull(out_type_, in->length(), pool_).try_value(out);
            }
            const auto &dict_array = GetDictionaryArray(in);

            using Builder = typename TypeTraits<T>::BuilderType;
            Builder builder(out_type_, pool_);
            TURBO_RETURN_NOT_OK(builder.resize(dict_array.indices()->length()));
            const auto &decimal_type(turbo::checked_cast<const DecimalType &>(*out_type_));
            int32_t out_precision = decimal_type.precision();
            int32_t out_scale = decimal_type.scale();

            auto visit_valid = [&](std::string_view repr) {
                value_type value;
                int32_t precision, scale;
                TURBO_RETURN_NOT_OK(TypeTraits<T>::BuilderType::value_type::from_string(
                        repr, &value, &precision, &scale));
                if (precision > out_precision) {
                    return GenericConversionError(*out_type_, ": ", repr, " requires precision ",
                                                  precision);
                }
                if (scale != out_scale) {
                    auto result = value.Rescale(scale, out_scale);
                    if (TURBO_UNLIKELY(!result.ok())) {
                        return GenericConversionError(*out_type_, ": ", repr, " requires scale ",
                                                      scale);
                    } else {
                        value = result.move_value_unsafe();
                    }
                }
                builder.unsafe_append(value);
                return turbo::OkStatus();
            };

            auto visit_null = [&builder]() {
                builder.unsafe_append_null();
                return turbo::OkStatus();
            };

            TURBO_RETURN_NOT_OK(VisitDictionaryEntries(dict_array, visit_valid, visit_null));
            return builder.finish(out);
        }
    };

    template<typename DateTimeType>
    class DateTimeConverter : public PrimitiveConverter {
    public:
        DateTimeConverter(MemoryPool *pool, const std::shared_ptr<DataType> &type)
                : PrimitiveConverter(pool, type), converter_(pool, repr_type()) {}

        turbo::Status Convert(const std::shared_ptr<Array> &in, std::shared_ptr<Array> *out) override {
            if (in->type_id() == Type::NA) {
                return MakeArrayOfNull(out_type_, in->length(), pool_).try_value(out);
            }

            std::shared_ptr<Array> repr;
            TURBO_RETURN_NOT_OK(converter_.Convert(in, &repr));

            auto out_data = repr->data()->copy();
            out_data->type = out_type_;
            *out = make_array(out_data);

            return turbo::OkStatus();
        }

    private:
        using ReprType = typename CTypeTraits<typename DateTimeType::c_type>::ArrowType;

        static std::shared_ptr<DataType> repr_type() {
            return TypeTraits<ReprType>::type_singleton();
        }

        NumericConverter<ReprType> converter_;
    };

    template<typename T>
    class BinaryConverter : public PrimitiveConverter {
    public:
        using PrimitiveConverter::PrimitiveConverter;

        turbo::Status Convert(const std::shared_ptr<Array> &in, std::shared_ptr<Array> *out) override {
            if (in->type_id() == Type::NA) {
                return MakeArrayOfNull(out_type_, in->length(), pool_).try_value(out);
            }
            const auto &dict_array = GetDictionaryArray(in);

            using Builder = typename TypeTraits<T>::BuilderType;
            Builder builder(out_type_, pool_);
            TURBO_RETURN_NOT_OK(builder.resize(dict_array.indices()->length()));

            // TODO(bkietz) this can be computed during parsing at low cost
            int64_t data_length = 0;
            auto visit_lengths_valid = [&](std::string_view value) {
                data_length += value.size();
                return turbo::OkStatus();
            };

            auto visit_lengths_null = [&]() {
                // no-op
                return turbo::OkStatus();
            };

            TURBO_RETURN_NOT_OK(
                    VisitDictionaryEntries(dict_array, visit_lengths_valid, visit_lengths_null));
            TURBO_RETURN_NOT_OK(builder.reserve_data(data_length));

            auto visit_valid = [&](std::string_view value) {
                builder.unsafe_append(value);
                return turbo::OkStatus();
            };

            auto visit_null = [&]() {
                builder.unsafe_append_null();
                return turbo::OkStatus();
            };

            TURBO_RETURN_NOT_OK(VisitDictionaryEntries(dict_array, visit_valid, visit_null));
            return builder.finish(out);
        }
    };

    turbo::Status make_converter(const std::shared_ptr<DataType> &out_type, MemoryPool *pool,
                         std::shared_ptr<Converter> *out) {
        switch (out_type->id()) {
#define CONVERTER_CASE(TYPE_ID, CONVERTER_TYPE)              \
  case TYPE_ID:                                              \
    *out = std::make_shared<CONVERTER_TYPE>(pool, out_type); \
    break
            CONVERTER_CASE(Type::NA, NullConverter);
            CONVERTER_CASE(Type::BOOL, BooleanConverter);
            CONVERTER_CASE(Type::INT8, NumericConverter<Int8Type>);
            CONVERTER_CASE(Type::INT16, NumericConverter<Int16Type>);
            CONVERTER_CASE(Type::INT32, NumericConverter<Int32Type>);
            CONVERTER_CASE(Type::INT64, NumericConverter<Int64Type>);
            CONVERTER_CASE(Type::UINT8, NumericConverter<UInt8Type>);
            CONVERTER_CASE(Type::UINT16, NumericConverter<UInt16Type>);
            CONVERTER_CASE(Type::UINT32, NumericConverter<UInt32Type>);
            CONVERTER_CASE(Type::UINT64, NumericConverter<UInt64Type>);
            CONVERTER_CASE(Type::FP32, NumericConverter<Fp32Type>);
            CONVERTER_CASE(Type::FP64, NumericConverter<Fp64Type>);
            CONVERTER_CASE(Type::TIMESTAMP, NumericConverter<TimestampType>);
            CONVERTER_CASE(Type::TIME32, DateTimeConverter<Time32Type>);
            CONVERTER_CASE(Type::TIME64, DateTimeConverter<Time64Type>);
            CONVERTER_CASE(Type::DATE32, DateTimeConverter<Date32Type>);
            CONVERTER_CASE(Type::DATE64, DateTimeConverter<Date64Type>);
            CONVERTER_CASE(Type::BINARY, BinaryConverter<BinaryType>);
            CONVERTER_CASE(Type::STRING, BinaryConverter<StringType>);
            CONVERTER_CASE(Type::LARGE_BINARY, BinaryConverter<LargeBinaryType>);
            CONVERTER_CASE(Type::LARGE_STRING, BinaryConverter<LargeStringType>);
            CONVERTER_CASE(Type::BINARY_VIEW, BinaryConverter<BinaryViewType>);
            CONVERTER_CASE(Type::STRING_VIEW, BinaryConverter<StringViewType>);
            CONVERTER_CASE(Type::DECIMAL128, DecimalConverter<Decimal128Type>);
            CONVERTER_CASE(Type::DECIMAL256, DecimalConverter<Decimal256Type>);
            default:
                return turbo::unimplemented_error("JSON conversion to ", *out_type,
                                              " is not supported");
#undef CONVERTER_CASE
        }
        return turbo::OkStatus();
    }

    const PromotionGraph *get_promotion_graph() {
        static struct : PromotionGraph {
            std::shared_ptr<Field> Null(const std::string &name) const override {
                return field(name, null(), true, Kind::Tag(Kind::kNull));
            }

            std::shared_ptr<DataType> Infer(
                    const std::shared_ptr<Field> &unexpected_field) const override {
                auto kind = Kind::FromTag(unexpected_field->metadata());
                switch (kind) {
                    case Kind::kNull:
                        return null();

                    case Kind::kBoolean:
                        return boolean();

                    case Kind::kNumber:
                        return int64();

                    case Kind::kString:
                        return timestamp(TimeUnit::SECOND);

                    case Kind::kArray: {
                        const auto &type = turbo::checked_cast<const ListType &>(*unexpected_field->type());
                        auto value_field = type.value_field();
                        return list(value_field->with_type(Infer(value_field)));
                    }
                    case Kind::kObject: {
                        auto fields = unexpected_field->type()->fields();
                        for (auto &field: fields) {
                            field = field->with_type(Infer(field));
                        }
                        return STRUCT(std::move(fields));
                    }
                    default:
                        return nullptr;
                }
            }

            std::shared_ptr<DataType> Promote(
                    const std::shared_ptr<DataType> &failed,
                    const std::shared_ptr<Field> &unexpected_field) const override {
                switch (failed->id()) {
                    case Type::NA:
                        return Infer(unexpected_field);

                    case Type::TIMESTAMP:
                        return utf8();

                    case Type::INT64:
                        return float64();

                    default:
                        return nullptr;
                }
            }
        } impl;

        return &impl;
    }

}  // namespace nebula::json
